package router

import (
	"fmt"
	"github.com/PaulXu-cn/goeval"
	"github.com/dop251/goja"
	"github.com/gin-gonic/gin"
	"html/template"
	"io"
	"net/http"
	"os/exec"
	"regexp"
)

type CodeInjectionHandler struct {
}

func NewCodeInjectionHandler() *CodeInjectionHandler {
	return &CodeInjectionHandler{}
}

func (handler *CodeInjectionHandler) RegisterRoutes(r *gin.RouterGroup) {
	r.POST("/eval", handler.Eval)
	r.POST("/run_jscode", handler.RunJsCode)
	r.GET("/text_template", handler.TextTemplate)
}

func (handler *CodeInjectionHandler) Eval(ctx *gin.Context) {
	// 从body中读取要执行的代码
	body, err := io.ReadAll(ctx.Request.Body)
	if err != nil {
		ctx.String(http.StatusInternalServerError, "Error: Failed to read request body")
		return
	}

	// 获取需要导入的包
	imports := ctx.QueryArray("imports")

	// 检查导包以及代码是否安全
	if !isSafeCode(string(body), imports) {
		ctx.String(http.StatusBadRequest, "Error: invaild code")
		return
	}

	re, err := goeval.Eval("", string(body),  imports...)
	if err != nil {
		ctx.String(http.StatusBadRequest, "Error: %v", err)
	} else {
		ctx.String(http.StatusOK, string(re))
	}
}

func (handler *CodeInjectionHandler) RunJsCode(ctx *gin.Context) {
	body, err := io.ReadAll(ctx.Request.Body)
	if err != nil {
		ctx.String(http.StatusInternalServerError, "Error: Failed to read request body")
		return
	}

	// 创建一个新的 JavaScript 运行时
	vm := goja.New()
	vm.Set("runCmd", runCmd)

	//未做安全限制，存在代码注入漏洞
	res, err := vm.RunString(string(body))

	if err != nil {
		ctx.String(http.StatusBadRequest, "Error: %v", err)
	} else {
		ctx.String(http.StatusOK, "%v", res)
	}
}

func (handler *CodeInjectionHandler) TextTemplate(ctx *gin.Context) {
	message := ctx.Query("message")

	// 解析模板文件
	t, err := template.ParseFiles("template.html")
	if err != nil {
		ctx.String(http.StatusInternalServerError, "Error: %v", err)
		return
	}

	// 定义数据
	data := struct {
		Title   string
		Heading string
		Message string
	}{
		Title:   "My Page",
		Heading: "Welcome to My Page",
		Message: message,	//未做安全限制，存在代码注入漏洞
	}

	err = t.Execute(ctx.Writer, data)
	if err != nil {
		ctx.String(http.StatusInternalServerError, "Error: %v", err)
		return
	}
}

func runCmd(cmd string) string {
	command := exec.Command("bash", "-c", cmd)

	// 使用 CombinedOutput 方法执行命令并捕获标准输出和标准错误
	output, err := command.CombinedOutput()
	if err != nil {
		return fmt.Sprintf("Error: %v\n", err)
	}
	return string(output)
}

func isSafeCode(code string, imports []string) bool {
	// 定义允许的包和函数
	allowedPackages := map[string]bool{"fmt": true}
	allowedFunctions := map[string]bool{"Println": true, "Sprintf": true}

	// 检查导入的包
	for _, pkg := range imports {
		if !allowedPackages[pkg] {
			return false
		}
	}

	//提取代码中的函数调用
	calls := extractFunctionCalls(code)
	// 检查调用的函数
	for _, call := range calls {
		if !allowedFunctions[call] {
			return false
		}
	}

	return true
}

func extractFunctionCalls(code string) []string {
	// 正则表达式匹配函数调用，例如: funcName(...
	re := regexp.MustCompile(`\b([a-zA-Z_][a-zA-Z0-9_]*)\s*\(`)

	// 找到所有匹配项
	matches := re.FindAllStringSubmatch(code, -1)

	// 提取函数名
	var functionCalls []string
	for _, match := range matches {
		if len(match) > 1 {
			functionCalls = append(functionCalls, match[1])
		}
	}

	return functionCalls
}